{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45ff14a2",
   "metadata": {},
   "source": [
    "## What可视化\n",
    "\n",
    "2D Grand CAM可视化模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d998c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 获得视频教程\n",
    "from onekey_algo.custom.Manager import onekey_show\n",
    "onekey_show('What模型可视化')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7fa11b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'\n",
    "import monai\n",
    "from glob import glob\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "mydir = r'C:\\Users\\onekey\\Desktop\\demo\\20250831\\peri5'\n",
    "samples = [os.path.join(mydir, f) for f in os.listdir(mydir) if f.endswith('.jpg') or f.endswith('.png')]\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f9a9194",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onekey_algo.custom.components.comp2 import extract, init_from_model, init_from_onekey\n",
    "\n",
    "model, transformer, device = init_from_onekey(r'models/peri5/resnet18/viz/')\n",
    "for n, m in model.named_modules():\n",
    "    print('Feature name:', n, \"|| Module:\", m)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e46336a1",
   "metadata": {},
   "source": [
    "## 可视化卷积层\n",
    "\n",
    "`Feature name:` 之后的名称为要可视化的层，例如`layer4.2.conv3`, 一般深度学习特征提取最后一层卷积层\n",
    "\n",
    "** 注意 ** : 可视化的层，一定为带有`conv`的卷积层，而且一般是最后一层。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d9ad6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_layer = \"layer4.1.conv2\"\n",
    "gradcam = monai.visualize.GradCAM(nn_module=model, target_layers=target_layer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39c9708b",
   "metadata": {},
   "source": [
    "## 打印可视化界面"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9eb2b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onekey_algo.datasets.image_loader import default_loader\n",
    "from onekey_algo.custom.components.comp2 import show_cam_on_image\n",
    "import torch\n",
    "\n",
    "viz_dir = r'models/peri5/resnet18/Grad-CAM' \n",
    "os.makedirs(viz_dir, exist_ok=True)\n",
    "for sample in samples:\n",
    "    if os.path.exists(os.path.join(viz_dir, os.path.basename(sample))):\n",
    "        continue\n",
    "    img = default_loader(sample)\n",
    "    sample_ = transformer(img)\n",
    "    sample_  = sample_.view(1, *sample_.size()).to(device)\n",
    "    res_cam = gradcam(x=sample_, class_idx=None)\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(12, 4), facecolor='white')\n",
    "#     axes[0].imshow(-res_cam[0][0].cpu(), cmap='jet')\n",
    "    axes[0].imshow(img.resize(sample_.size()[2:]))\n",
    "    axes[0].axis('off')\n",
    "#     plt.savefig(f\"viz/{os.path.basename(sample).replace('.png', '_se.png')}\", bbox_inches = 'tight')\n",
    "#     plt.show()\n",
    "#     plt.figure(figsize=(10, 10))\n",
    "#     plt.axis('off')\n",
    "    imshow = axes[1].imshow(-res_cam[0][0].cpu(),cmap='jet')\n",
    "    axes[1].axis('off')\n",
    "    imshow = axes[2].imshow(show_cam_on_image(img.resize(sample_.size()[2:]), -res_cam[0][0].cpu(), use_rgb=True, reverse=False), \n",
    "                            cmap='jet')\n",
    "    axes[2].axis('off')\n",
    "    cax = fig.add_axes([0.92, 0.17, 0.02, axes[2].get_position().height]) \n",
    "    plt.colorbar(imshow, cax=cax)\n",
    "    plt.savefig(f'{viz_dir}/{os.path.basename(sample).replace(\".npy\", \".png\")}', bbox_inches = 'tight')\n",
    "#     plt.show()\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5048302a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
